use anyhow::{Context, Result};
use cudarc::driver::{CudaDevice, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, DriverError};

use std::ffi::c_void;

// 声明外部CUDA函数
extern "C" {
    fn launch_vector_add(
        a: *const f32,
        b: *const f32,
        result: *mut f32,
        n: usize,
        stream: *mut std::ffi::c_void,
    );
}

fn main() -> Result<()> {
    // 初始化CUDA设备
    let device = CudaDevice::new(0)?;

    // 准备数据
    let n = 1024;
    let a: Vec<f32> = (0..n).map(|x| x as f32).collect();
    let b: Vec<f32> = (0..n).map(|x| (x * 2) as f32).collect();
    let mut result = vec![0.0f32; n];

    // 分配设备内存
    let d_a: CudaSlice<f32> = device.htod_sync_copy(&a)?;
    let d_b: CudaSlice<f32> = device.htod_sync_copy(&b)?;
    let mut d_result: CudaSlice<f32> = device.alloc_zeros(n)?;

    // 创建CUDA流
    let stream = device.fork_default_stream()?;

    // 获取原始指针
    let a_ptr = *d_a.device_ptr() as *const f32;
    let b_ptr = *d_b.device_ptr() as *const f32;
    let result_ptr = *d_result.device_ptr_mut() as *mut f32;
    let stream_ptr = stream.stream as *mut c_void;

    // 启动内核
    unsafe {
        launch_vector_add(a_ptr, b_ptr, result_ptr, n, stream_ptr);
    }

    // 等待流完成
    device.wait_for(&stream)?;

    // 复制结果回主机
    device.dtoh_sync_copy_into(&d_result, &mut result)?;
       // 验证结果
    for i in 0..n {
        assert_eq!(result[i], a[i] + b[i]);
    }

    println!("CUDA vector add executed successfully!");
    Ok(())
}
